package com.custom.extra;


import com.custom.MySplit;
import lombok.val;
import org.apache.flink.connector.base.source.reader.RecordsBySplits;
import org.apache.flink.connector.base.source.reader.RecordsWithSplitIds;
import org.apache.flink.connector.base.source.reader.splitreader.SplitReader;
import org.apache.flink.connector.base.source.reader.splitreader.SplitsChange;

import java.io.IOException;
import java.util.ArrayDeque;
import java.util.HashSet;
import java.util.Queue;
import java.util.Set;

/**
 * 分片读取器
 */
public class MySplitReader implements SplitReader<Long, MySplit> {

    private final Queue<MySplit> splits = new ArrayDeque<>();

    private final Set<String> finishedSplits = new HashSet<>();

    @Override
    public RecordsWithSplitIds<Long> fetch() throws IOException {
        val split = splits.poll();
        RecordsBySplits.Builder<Long> builder = new RecordsBySplits.Builder<>();
        if (split != null) {
            while (split.hasNext()) {
                builder.add(split.splitId(), split.getNext());
            }
            finishedSplits.add(split.splitId());
        }

        builder.addFinishedSplits(finishedSplits);
        return builder.build();
    }

    @Override
    public void handleSplitsChanges(SplitsChange<MySplit> splitsChanges) {
        splits.addAll(splitsChanges.splits());
    }

    @Override
    public void wakeUp() {

    }

    @Override
    public void close() throws Exception {

    }
}
